#!/usr/bin/env python
# encoding: utf-8
"""
mleRelatedness.py

Created by Brant Faircloth on 2007-12-11.
Copyright (c) 2007 Brant C. Faircloth. All rights reserved.

Pairwise relatedness estimation using null allele frequency estimate and downhill
Simplex algorithm for optimization.

"""

import csv,pdb,sqlite3,scipy.optimize
from numpy import *
from copy import deepcopy

# the following frequencies are output from FreeNA and converted to dict format
# you could estimate them using something else.  For details on FreeNA, see
# Chapuis and Estoup, 2007, MBE

# allele frequencies adjusted for null, as dict
nullFreq = {}
# the actual NULL allele frequencies
null = {}
# input file
my_input 	= 'someFile.csv'
# output file
my_output 	= 'someOtherFile.csv'


allaicWeights = []

class excelSingleSpace:
    """class for csv module to work correctly"""
    delimiter = ','
    quotechar = '"'
    escapechar = None
    doublequote = True
    skipinitialspace = False
    lineterminator = '\r'

def read(file):
    handle  = open(file,'rU')
    c       = csv.reader(handle, delimiter=',', dialect = excelSingleSpace)
    # skip remaining header
    count = {}
    header = c.next()
    row = header[2:]
    lociCount= range(len(row)/2)
    loci = range(0,len(row),2)
    for l in lociCount:
        count[l] = {}
    #pdb.set_trace()
    for line in c:
        line = line[2:]
        for allele in loci:
            if line[allele]=='*':
                pass
            elif line[allele] not in count[loci.index(allele)].keys():
                count[loci.index(allele)][line[allele]] = 1
            else:
                count[loci.index(allele)][line[allele]] +=1
            if line[allele+1]=='*':
                pass
            elif line[allele+1] not in count[loci.index(allele)].keys():
                count[loci.index(allele)][line[allele+1]] = 1
            else:
                count[loci.index(allele)][line[allele+1]] +=1
    handle.close()
    freq = {}
    sort = {}
    #pdb.set_trace()
    for locus in count.keys():
        lKeys = count[locus].keys()
        # sort the keys to place loci in order
        lKeys.sort()
        aList = []
        sort[locus] = []
        for alleles in lKeys:
            aList.append(count[locus][alleles])
            sort[locus].append(alleles)
        aArray = array(aList)
        aFreq = aArray/float(sum(aArray))
        freq[locus] = aFreq
    return header,freq,sort
    #pdb.set_trace()


def reformat(sort,file):
    handle  = open(file,'rU')
    c       = csv.reader(handle, delimiter=',', dialect = excelSingleSpace)
    # skip remaining header
    records = {}
    row = c.next()
    for row in c:
        group = row[1]
        line = row[2:]
        loci = range(0,len(line),2)
        genotype = []
        for allele in loci:
            if line[allele] == '*':
                genotype.append(None)
                genotype.append(None)
            else:
                genotype.append(sort[loci.index(allele)].index(line[allele]) + 1)
                genotype.append(sort[loci.index(allele)].index(line[allele + 1]) + 1)
        records[row[0]] = [[group],genotype]
    #pdb.set_trace()
    return records

def dbase(header,records,c):
    '''
    keep the data (transformed) in the database to make dealing w/
    groups easier than it would be by sorting.  Obviously requires sqlite3.
    '''
    # drop tables if they exist
    try:
        # if previous tables exist, drop them
        # TODO: fix createDbase() to drop tables safely
        c.execute('''DROP TABLE genotypes''')
        # commit the change
        c.commit()
    except:
        pass
    header[0] += " text"
    header[1] += " int"
    for item in range(2,len(header),1):
        header[item] += " int"  
    query = '''CREATE TABLE genotypes (''' + ','.join(header) + ''')'''
    columns = []
    for item in range(len(header)):
        columns.append('?')
    #pdb.set_trace()
    c.execute(query)
    c.commit()
    recordKeys = records.keys()
    recordKeys.sort()
    #pdb.set_trace()
    for bird in recordKeys:
        data = (bird,)
        if records[bird][0][0] =='':
            data += (None,)
        else:
            data += (records[bird][0][0],)
        for item in records[bird][1]:
            data += (item,)
        query = '''INSERT INTO genotypes VALUES (''' + ','.join(columns) + ''')'''
        #pdb.set_trace()
        c.execute(query,data)
    c.commit()


def objFunc(k_coeff,values, p = False):
    # if values are unreasonable, kick back a huge number
    if k_coeff[0] < 0. or k_coeff[1] < 0.:
        return Inf
    elif k_coeff[0] > 1. or k_coeff[1] > 1.:
        return Inf
    elif sum(k_coeff) > 1.:
        return Inf
    else:
        k0          = 1 - sum(k_coeff)
        k_coeff     = hstack((k_coeff,k0))
        # to check for optim. errors
        if p: print k_coeff
        # mult k_coeff by array of values for each locus
        temp        = k_coeff * values
        # sum across rows
        rowSum      = sum(temp,axis=1)
        truthValues = rowSum == 0
        # if there are zeros being multiplied, return big number
        if truthValues.any():
            logLike = Inf
        else:
            logLike = -log(product(rowSum))
        return logLike


def relatedness(freqs,kids,cw):
    kidsSort = kids.keys()
    kidsSort.sort()
    kIter = 0
    header = deepcopy(kidsSort)
    cw.writerow(['']+header)
    for init in kidsSort:
        #pdb.set_trace()
        base = kids[init]
        kIter += 1
        outRow = []
        for compare in kidsSort[0:kIter]:
            other = kids[compare]
            if init == compare:
                r = ''
            else:
                values = zeros((14,3))
                locus = 0
                for alleles in range(0,len(base),2):
                    allele1,allele2 = alleles,alleles+2
                    cDict = {}
                    bloc = base[allele1:allele2] 
                    cloc = other[allele1:allele2]
                    concat = bloc + cloc
                    for a in concat:
                        if a in cDict.keys():
                            cDict[a] += 1
                        else:
                            if a:
                                cDict[a] = 1
                            else:
                                #pdb.set_trace()
                                cDict[NaN] = NaN
                    cDictV,cDictK = array(cDict.values()), array(cDict.keys())
                    ind = lexsort(keys = (cDictK,cDictV))
                    cDictV,cDictK = cDictV[ind], cDictK[ind]
                    ####################################
                    # deal with noneType
                    ####################################
                    if not all(isfinite(cDictV)):
                        #pdb.set_trace()
                        values[locus][0] = NaN
                        values[locus][1] = NaN
                        values[locus][2] = NaN
                    ####################################
                    #homozygous; 2 matches
                    ####################################
                    elif all(cDictV == array([4])):
                        p = freqs[locus][cDictK[0]-1]
                        values[locus][0] = p**2 + (0. + 0. + 2 * p * null[locus])
                        values[locus][1] = p**3 + (2 * p**2 * null[locus]) + (p * null[locus] * (p + null[locus]))
                        values[locus][2] = p**4 + (4 * p**3 * null[locus]) + (4 * p**2 * null[locus]**2)
                    ####################################
                    #homozygous; 0 matches
                    ####################################
                    elif all(cDictV == array([2,2])) and (concat[0] == concat[1] and concat[2] == concat[3]):
                        p = freqs[locus][cDictK[1]-1]
                        j = freqs[locus][cDictK[0]-1]
                        # DEFAULT: values slot 1 = 0
						values[locus][1] = 0 + (0 + 0 + p * j * null[locus])
                        values[locus][2] = p**2 * j**2 + (2 * p**2 * j * null[locus]) + (2 * p * j**2 * null[locus]) + (4 * p * j * null[locus]**2)
                    ####################################
                    #homozygous-heterozygous; 1 match
                    ####################################
                    elif all(cDictV == array([1,3])):
                        p = freqs[locus][cDictK[1]-1]
                        j = freqs[locus][cDictK[0]-1]
                        #values slot 1 = 0
                        values[locus][1] = p**2 * j + (p * j * null[locus])
                        values[locus][2] = 2 * p**3 * j + (4 * p**2 * j * null[locus])
                    ####################################
                    #homozygous-heterozygous; 0 matches
                    ####################################
                    elif all(cDictV == array([1,1,2])) and (concat[0] == concat[1] or concat[2] == concat[3]):
                        p = freqs[locus][cDictK[2]-1]
                        j = freqs[locus][cDictK[1]-1]
                        k = freqs[locus][cDictK[0]-1]
                        #values slot 0 = 0
                        #values slot 1 = 0
                        values[locus][2] = 2 * p**2 * j * k + (4 * p * j * k * null[locus])
                    ####################################
                    #heterozygote-heterozygote; 2 matches
                    ####################################
                    elif all(cDictV == array([2,2])) and (concat[0] != concat[1] and concat[2] != concat[3]):
                        p = freqs[locus][cDictK[1]-1]
                        j = freqs[locus][cDictK[0]-1]
                        values[locus][0] = 2 * p * j
                        values[locus][1] = (p * j) * (p + j)
                        values[locus][2] = 4 * p**2 * j**2
                    ####################################
                    #heterozygote-heterozygote; 1 match
                    ####################################
                    elif all(cDictV == array([1,1,2])) and (concat[0] != concat[1] and concat[2] != concat[3]):
                        p = freqs[locus][cDictK[2]-1]
                        j = freqs[locus][cDictK[1]-1]
                        k = freqs[locus][cDictK[0]-1]
                        #values slot 0 = 0
                        values[locus][1] = p * j * k
                        values[locus][2] = 4 * p**2 * j * k
                    ####################################
                    #heterozygote-heterozygote; 0 matches
                    ####################################
                    elif all(cDictV == array([1,1,1,1])):
                        p = freqs[locus][cDictK[3]-1]
                        j = freqs[locus][cDictK[2]-1]
                        k = freqs[locus][cDictK[1]-1]
                        l = freqs[locus][cDictK[0]-1]
                        #values slot 0 = 0
                        #values slot 1 = 0
                        values[locus][2] = 4 * p * j * k * l
                    ####################################
                    #catch any weirdness
                    ####################################
                    else:
                        print 'Skipped a pair because loci could not be matched.'
                        pdb.set_trace()
                    locus += 1
                nans = isfinite(values)
                # reindex array by Trues (remove NaNs)
                values = values[nans[:,0]]
                locusCount = len(values)
                # starting values for optimization - start in middle
                x0 = (0.33,0.33)
				# do optimization by simplex algorithm
                opt = scipy.optimize.fmin(objFunc,x0,args=(values,),maxiter=1000,maxfun=1000,full_output=1,disp=0)
                k_coeff, warn = opt[0], opt[-1]
                if warn == 0:
                    r = 2*(k_coeff[1]/4. + k_coeff[0]/2.)
                    r = round(r,4)
                    if r == 0:r=0.0000
                else:
                    print 'Optimization did not converge for:', init,compare
                    r = '\N'
                #print '----------'
            outRow.append(r)
        outRow = [init] + outRow
        cw.writerow(outRow)

def groups(c,freq=nullFreq,null=null):
    # get the distinct groups in the data set
    c.execute('SELECT DISTINCT(groups) FROM genotypes WHERE groups IS NOT NULL ORDER BY groups')
    groups = c.fetchall()
    # outfile name:
	outfile = open(my_output,'w')
    cw = csv.writer(outfile, dialect = excelSingleSpace)
    for group in groups:
        cw.writerow([group])
        c.execute('SELECT DISTINCT(bird_id) FROM genotypes WHERE groups= ? ORDER BY bird_id', group)
        birds = c.fetchall()
        #pdb.set_trace()
        groupRecords = {}
        print '#####################'
        print '\t',group[0]
        print '#####################'
        for bird in birds:
            c.execute('SELECT * FROM genotypes WHERE bird_id = ?', bird)
            data = c.fetchall()
            groupRecords[bird[0]] = list(data[0][2:])
        relatedness(freq,groupRecords,cw)
        cw.writerow([''])
        #pdb.set_trace()
    outfile.close()

def main():
    c = sqlite3.connect(':memory:')
    cur = c.cursor()
    header, trash, sort = read(my_input)
    records = reformat(sort,my_input)
    dbase(header,records,c)
    #pdb.set_trace()
    if nullFreq and null:
        groups(cur,nullFreq,null)
    else:
        groups(cur)
    #pdb.set_trace()

if __name__ == '__main__':
    main()